"""
手动指定gpu 可以在模型很大 一块内存不够时 把模型分在几个不同的gpu上 
但是本质仍是穿行的

"""
import tensorflow as tf 
from tensorflow import keras

tf.debugging.set_log_device_placement(True)  # 显示模型参数在哪块gpu上

gpus = tf.config.experimental.list_physical_devices('GPU')
for gpu in gpus:
    tf.config.experimental.set_memory_growth(True,gpu)  # 设置内存自增长

logical_gpus = tf.config.experimental.list_logical_devices('GPU')  # 操作逻辑设备 不要直接操作物理设备

model = keras.models.Sequential()

with tf.device(logical_gpus[0].name):  # 第一块gpu分100层
    model.add(keras.layers.Dense(100,activation='relu',input_shape=(8,)))
    [model.add(keras.layers.Dense(100,activation='relu')) for _ in range(99)]

with tf.device(logical_gpus[1].name):  # 第二块gpu分100层
    [model.add(keras.layers.Dense(100,activation='relu')) for _ in range(100)]

with tf.device(logical_gpus[2].name):  # 第三块gpu分100层
    [model.add(keras.layers.Dense(100,activation='relu')) for _ in range(100)]

with tf.device(logical_gpus[3].name):  # 第四块gpu分100层
    [model.add(keras.layers.Dense(100,activation='relu')) for _ in range(9)]
    model.add(keras.layers.Dense(1,activation='softmax'))

model.summary()  # 一个大的模型被分到了四块gpu上


